import pymysql
import xlwt


def export_excel(db_name, table_name):
    # 连接数据库
    conn = pymysql.connect(
        host="118.24.3.40",
        user="jxz",
        password="123456",
        db="",
        port=3306,
        charset="utf8",
        autocommit=True
    )
    cur = conn.cursor(pymysql.cursors.DictCursor)  # 建立游标
    # 拼接库名表名
    new_str = db_name + '.' + table_name
    test_connect_sql = 'select * from %s;' % new_str
    cur.execute(test_connect_sql)
    # 判断是否能连接库和表获取查询结果
    if cur.fetchall():
        query_sql = 'select * from %s;' % new_str
        cur.execute(query_sql)
        # 获取表数据
        all_data = cur.fetchall()
        # print(all_data)
        # 如果sql可以查询到数据
        if all_data:
            book = xlwt.Workbook()  # 先创建一个book
            sheet = book.add_sheet('sheet1')  # 创建一个sheet表
            for index, key in enumerate(all_data[0]):  # 写表头
                sheet.write(0, index, key)
            for row, item in enumerate(all_data, 1):  # 行的逻辑
                for col, value in enumerate(item.values()):  # 列的逻辑
                    sheet.write(row, col, value)
            # 以数据库表名命名
            book.save(table_name + '.xls')
        else:
            print("您输入的库名和表名错误，sql无法正常查询，请检查！！")
    cur.close()
    conn.close()


if __name__ == '__main__':
    db_name = input("输入你要传入的数据库名").strip()
    table_name = input("输入你要传入的表名").strip()
    export_excel(db_name, table_name)
